package github.sf.fw.tools.security;


import javax.crypto.Cipher;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.KeySpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * 非对称加密工具类
 */
public abstract class RSAEncrypt {
    private RSAEncrypt() {
    }

    /**
     * 字符串公钥
     */
    private String publicKey;
    /**
     * 字符串私钥
     */
    private String privateKey;
    /**
     * 二进制公钥
     */
    private byte[] publicKeyBytes;
    /**
     * 二进制私钥
     */
    private byte[] privateKeyBytes;
    private RSAPrivateKey rsaPrivateKey;
    private RSAPublicKey rsaPublicKey;
    /**
     * 是否私用私钥加密，公钥解密
     */
    private boolean encodeWithPrivateKey = false;
    /**
     * 是否使用自定义key,如果自己设置公钥和私钥解密，改选项必须为true
     */
    private boolean useKey = false;

    public String getPublicKey() {
        return publicKey;
    }

    /**
     * 可以通过修改encodeWithPrivateKey来使用私钥加密，公钥解密
     *
     * @param encodeWithPrivateKey 是否使用私钥加密
     */
    public void setEncodeWithPrivateKey(boolean encodeWithPrivateKey) {
        this.encodeWithPrivateKey = encodeWithPrivateKey;
    }

    public void setUseKey(boolean useKey) {
        this.useKey = useKey;
    }

    public void setPublicKey(String publicKey) {
        this.publicKey = publicKey;
    }

    public String getPrivateKey() {
        return privateKey;
    }

    public void setPrivateKey(String privateKey) {
        this.privateKey = privateKey;
    }

    public byte[] getPublicKeyBytes() {
        return publicKeyBytes;
    }

    public void setPublicKeyBytes(byte[] publicKeyBytes) {
        this.publicKeyBytes = publicKeyBytes;
    }

    public byte[] getPrivateKeyBytes() {
        return privateKeyBytes;
    }

    public void setPrivateKeyBytes(byte[] privateKeyBytes) {
        this.privateKeyBytes = privateKeyBytes;
    }


    public static RSAEncrypt getEncrypt() {
        return new RSAEncrypt() {
        };
    }

    public void genKeyPair() {
        KeyPairGenerator keyPairGenerator = null;
        try {
            keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e.getMessage(), e);
        }
        //初始化密钥对生成器，密钥大小为96-1024位
        keyPairGenerator.initialize(1024, new SecureRandom());
        KeyPair keyPair = keyPairGenerator.genKeyPair();
        rsaPrivateKey = (RSAPrivateKey) keyPair.getPrivate();
        rsaPublicKey = (RSAPublicKey) keyPair.getPublic();
        String privateKeyString = Base64.getEncoder().encodeToString(rsaPrivateKey.getEncoded());
        String publicKeyString = Base64.getEncoder().encodeToString(rsaPublicKey.getEncoded());
        this.privateKey = privateKeyString;
        this.publicKey = publicKeyString;
        this.publicKeyBytes = rsaPublicKey.getEncoded();
        this.privateKeyBytes = rsaPrivateKey.getEncoded();
    }

    /**
     * 加密(默认公钥加密)
     */
    public byte[] encrypt(String str) {
        checkAndSet();
        //RSA加密
        Cipher cipher = null;
        try {
            cipher = Cipher.getInstance("RSA");
            if (encodeWithPrivateKey) {
                if (useKey) {
                    RSAPrivateKey privkey = (RSAPrivateKey) KeyFactory.getInstance("RSA").generatePrivate(getKeySpec(Base64.getDecoder().decode(privateKey), 2));
                    cipher.init(Cipher.ENCRYPT_MODE, privkey);
                } else {
                    cipher.init(Cipher.ENCRYPT_MODE, rsaPrivateKey);
                }
            } else {
                if (useKey) {
                    RSAPublicKey pubkey = (RSAPublicKey) KeyFactory.getInstance("RSA").generatePublic(getKeySpec(Base64.getDecoder().decode(publicKey), 1));
                    cipher.init(Cipher.ENCRYPT_MODE, pubkey);
                } else {
                    cipher.init(Cipher.ENCRYPT_MODE, rsaPublicKey);
                }
            }
            return cipher.doFinal(str.getBytes(StandardCharsets.UTF_8));
        } catch (Exception e) {
            throw new SecurityException(e.getMessage(), e);
        }
    }

    public String encryptAsString(String str) {
        byte[] bytes = encrypt(str);
        return Base64.getEncoder().encodeToString(bytes);
    }

    /**
     * 解密(默认私钥解密)
     */
    public byte[] decrypt(String str) {
        checkAndSet();
        //RSA加密
        try {
            Cipher cipher = Cipher.getInstance("RSA");
            if (encodeWithPrivateKey) {
                if (useKey) {
                    RSAPublicKey pubkey = (RSAPublicKey) KeyFactory.getInstance("RSA").generatePublic(getKeySpec(Base64.getDecoder().decode(publicKey), 1));
                    cipher.init(Cipher.DECRYPT_MODE, pubkey);
                } else {
                    cipher.init(Cipher.DECRYPT_MODE, rsaPublicKey);
                }
            } else {
                if (useKey) {
                    RSAPrivateKey privkey = (RSAPrivateKey) KeyFactory.getInstance("RSA").generatePrivate(getKeySpec(Base64.getDecoder().decode(privateKey), 2));
                    cipher.init(Cipher.DECRYPT_MODE, privkey);
                } else {

                    cipher.init(Cipher.DECRYPT_MODE, rsaPrivateKey);
                }
            }
            byte[] decode = Base64.getDecoder().decode(str);
            return cipher.doFinal(decode);
        } catch (Exception e) {
            throw new SecurityException(e.getMessage(), e);
        }
    }

    public String decryptAsString(String str) {
        byte[] bytes = decrypt(str);
        return byteArrayAsString(bytes);
    }

    private String byteArrayAsString(byte[] bytes) {
        return new String(bytes, StandardCharsets.UTF_8);
    }

    private KeySpec getKeySpec(byte[] keyBytes, int num) {
        if (num == 1) {
            return new X509EncodedKeySpec(keyBytes);
        }
        return new PKCS8EncodedKeySpec(keyBytes);
    }

    private void checkAndSet() {
        if (rsaPublicKey == null || rsaPrivateKey == null) {
            useKey = true;
        } else {
            if (privateKey == null && publicKey == null) {
                throw new IllegalArgumentException("privateKey and public key can't be null in the same time.");
            }
            if (encodeWithPrivateKey && (privateKey == null || privateKey.isEmpty())) {
                throw new IllegalArgumentException("privateKey  can't be null when encrypt mode is encodeWithPrivateKey");
            }
            if (!encodeWithPrivateKey && (publicKey == null || publicKey.isEmpty())) {
                throw new IllegalArgumentException("publicKey  can't be null when encrypt mode is encodeWithPublicKey");
            }
        }
    }
}
